# 给定两个以 升序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。
#  定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。
#  请找到和最小的 k 个数对 (u1,v1), (u2,v2) ... (uk,vk) 。
#
#  示例 1:
# 输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
# 输出: [1,2],[1,4],[1,6]
# 解释: 返回序列中的前 3 对数：
#      [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
#
#  示例 2:
# 输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
# 输出: [1,1],[1,1]
# 解释: 返回序列中的前 2 对数：
#     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
#
#  示例 3:
# 输入: nums1 = [1,2], nums2 = [3], k = 3
# 输出: [1,3],[2,3]
# 解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]
import heapq
from typing import List


class Solution:
    def kSmallestPairs2(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        pass

    def kSmallestPairs1(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        heap = []
        for i in range(min(k, len(nums1))):
            for j in range(min(k, len(nums2))):
                if len(heap) != k or heap[0][0] <= -nums1[i] - nums2[j]:
                    if len(heap) == k:
                        heapq.heappop(heap)
                    heapq.heappush(heap, [-nums1[i] - nums2[j], -nums1[i], -nums2[j]])
        return list(map(lambda item: [-item[1], -item[2]], heap))

    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        return self.kSmallestPairs1(nums1, nums2, k)


if __name__ == "__main__":
    nums1, nums2 = [1, 7, 11], [2, 4, 6]
    k = 3
    print(Solution().kSmallestPairs(nums1, nums2, k))
